#pragma once
#include<iostream>

using std::pair;
using std::cout;
using std::endl;

enum Color{RED=1,BLACK};

template<class K,class V>
struct RBTreeNode
{
    pair<K,V> _kv;
    RBTreeNode<K,V>* _left;
    RBTreeNode<K,V>* _right;
    RBTreeNode<K,V>* _parent;

    Color _col;

    RBTreeNode(const pair<K,V>& kv)
    :_kv(kv)
    ,_left(nullptr)
    ,_right(nullptr)
    ,_parent(nullptr)
    ,_col(RED)
    {}
};

template<class K,class V>
class RBTree
{
    typedef RBTreeNode<K,V> Node;
public:
    RBTree():_root(nullptr){}

    bool Insert(const pair<K,V>& kv)
    {
        if(_root==nullptr)
        {
            _root=new Node(kv),_root->_col=BLACK;
            return true;
        }

        //1.根据kv.first，找位置
        Node* cur=_root,*parent=nullptr;
        while(cur)
        {
            if(kv.first>cur->_kv.first)
                parent=cur,cur=cur->_right;
            else if(kv.first<cur->_kv.first)
                parent=cur,cur=cur->_left;
            else return false;//说明有key值相同的节点
        }
        //2.将newnode插入树中
        cur=new Node(kv);
        if(kv.first>parent->_kv.first)
            parent->_right=cur,cur->_parent=parent;
        else 
            parent->_left=cur,cur->_parent=parent;
        //3.更新结点颜色
        while(parent && parent->_col==RED)
        {
            Node* grandparent=parent->_parent;//此时，grandparent必定不为空，且为黑色
            //判断parent是grandparent的哪棵子树
            if(parent==grandparent->_left)
            {
                Node* uncle=grandparent->_right;
                //判断uncle的状态
                if(uncle && uncle->_col==RED)
                {
                    parent->_col=uncle->_col=BLACK;
                    grandparent->_col=RED;//变色

                    cur=grandparent,parent=cur->_parent;//继续向上更新
                }
                //uncle不存在，或存在且为黑，需要旋转+变色
                else if(cur==parent->_left)//判断cur是parent的哪棵子树
                {
                    RotateR(grandparent);//右单旋
                    parent->_col=BLACK,grandparent->_col=RED;//变色
                }
                else if(cur==parent->_right)
                {
                    RotateL(parent);//先左旋
                    RotateR(grandparent);//再右旋
                    cur->_col=BLACK,grandparent->_col=RED;//变色
                    break;
                } 
            }
            else 
            {
                Node* uncle=grandparent->_left;
                //判断uncle的状态
                if(uncle && uncle->_col==RED)
                {
                    parent->_col=uncle->_col=BLACK;
                    grandparent->_col=RED;//变色

                    cur=grandparent,parent=cur->_parent;//继续向上更新
                }
                else if(cur==parent->_right)
                {
                    RotateL(grandparent);//左单旋
                    parent->_col=BLACK,grandparent->_col=RED;//变色
                }
                else if(cur==parent->_left)
                {
                    RotateR(parent);//先右旋
                    RotateL(grandparent);//再左旋
                    cur->_col=BLACK,grandparent->_col=RED;//变色
                    break;
                }
            }
        }
        _root->_col=BLACK;
        return true;
    }   

    void InOrder()
    {
        _InOrder(_root);
        cout<<endl;
    }

    bool Is_Balance()
    {
        //1.根节点黑色   
        if(_root->_col!=BLACK)return false;

        //找个基准值
        int black_size=0;
        Node* cur=_root;
        while(cur)
        {
            if(cur->_col==BLACK)black_size++;
            cur=cur->_left;
        }

        //2.不能有连续的红色节点   3.每条支路黑色结点个数相等
        return _Is_Red_Row(_root,0,black_size);
    }

private:
    void RotateL(Node* parent)
    {
        Node* grandparent=parent->_parent;
        Node* cur=parent->_right;
        Node* curleft=cur->_left;

        parent->_right=curleft;
        parent->_parent=cur;

        if(curleft)
            curleft->_parent=parent;

        cur->_left=parent;
        cur->_parent=grandparent;

        if(grandparent==nullptr)
            _root=cur;
        else 
        {
            if(grandparent->_left==parent)
                grandparent->_left=cur;
            else 
                grandparent->_right=cur;
        }
    }

    void RotateR(Node* parent)
    {
        Node* grandparent=parent->_parent;

        Node* cur=parent->_left;
        Node* curright=cur->_right;

        parent->_left=curright;
        parent->_parent=cur;

        if(curright)
            curright->_parent=parent;
        
        cur->_right=parent;
        cur->_parent=grandparent;

        if(grandparent==nullptr)
            _root=cur;
        else 
        {
            if(parent==grandparent->_left)
                grandparent->_left=cur;
            else 
                grandparent->_right=cur;
        }
    }

    void _InOrder(Node* root)
    {
        if(root==nullptr)
            return;
        _InOrder(root->_left);
        cout<<root->_kv.first<<" ";
        _InOrder(root->_right);
    }

    bool _Is_Red_Row(Node* root,int black_number,int standard)
    {
        if(root==nullptr)
        {
            if(black_number!=standard)return false;

            return true;
        }

        if(root->_col==BLACK)black_number++;

        if(root->_col==RED && root->_parent->_col==RED)
            return false;

        return _Is_Red_Row(root->_left,black_number,standard) && _Is_Red_Row(root->_right,black_number,standard);
    }

private:
    Node* _root;
};